#include <iostream>
#include <vector>
#include <math.h>
#include <assert.h>

#include "QuadElement.hpp"
#include "BernsteinBasis.hpp"

QuadElement::QuadElement( int element_index,
    const class MeshInfo_2d &mesh2d,
    std::vector<double> ctrlPts,
    int cpdim,
    std::vector<NURBS_T::BezierElem*> seg_x,
    std::vector<NURBS_T::BezierElem*> seg_y )
{
  en_index   = element_index;
  en_index_x = (en_index + 1) % mesh2d.nElem_x;
  en_index_y = (en_index + 1 - en_index_x) / mesh2d.nElem_x;
  en_index_x = en_index - 1;

  en_hx = mesh2d.hx[en_index_x];
  en_hy = mesh2d.hy[en_index_y];

  qp_x = mesh2d.qp_x;
  wq_x = mesh2d.wq_x;
  qp_y = mesh2d.qp_y;
  wq_y = mesh2d.wq_y;

  enqp_x = mesh2d.nqp_x;
  enqp_y = mesh2d.nqp_y;
  sDegree = mesh2d.sDegree;
  tDegree = mesh2d.tDegree;

  enLocBas = mesh2d.nLocBas;
  localIEN = new int [enLocBas];
  for(int i=0; i<enLocBas; i++)
    localIEN[i] = mesh2d.IEN[en_index][i];

  ctrl_x = new double [enLocBas];
  ctrl_y = new double [enLocBas];
  ctrl_w = new double [enLocBas];
  R      = new double [enLocBas];
  dR_dx  = new double [enLocBas];
  dR_dy  = new double [enLocBas];
  d2R_dx = new double [enLocBas];
  d2R_dy = new double [enLocBas];
  d2R_dxdy = new double [enLocBas];
  
  for(int i=0; i<enLocBas; i++)
  {
    int index = localIEN[i];
    ctrl_x[i] = ctrlPts[index*cpdim+0];
    ctrl_y[i] = ctrlPts[index*cpdim+1];
    ctrl_w[i] = ctrlPts[index*cpdim+3];
    R[i]        = 0.0;
    dR_dx[i]    = 0.0;
    dR_dy[i]    = 0.0;
    d2R_dx[i]   = 0.0;
    d2R_dy[i]   = 0.0;
    d2R_dxdy[i] = 0.0;
  }

  Jacobi = 0.0;

  c_x = new double * [sDegree+1];
  c_y = new double * [tDegree+1];
  for(int i=0; i<sDegree+1; i++)
    c_x[i] = new double [sDegree+1];
  for(int i=0; i<tDegree+1; i++)
    c_y[i] = new double [tDegree+1];

  for(int i=0; i<sDegree+1; i++)
  {
    for(int j=0; j<sDegree+1; j++)
    {
      c_x[i][j] = seg_x[en_index_x]->coefficient(i,j);
    }
  }

  for(int i=0; i<tDegree+1; i++)
  {
    for(int j=0; j<tDegree+1; j++)
    {
      c_y[i][j] = seg_y[en_index_y]->coefficient(i,j);
    }
  }

};

/* Destructor */
QuadElement::~QuadElement()
{
  delete [] localIEN;
  delete [] ctrl_x;
  delete [] ctrl_y;
  delete [] ctrl_w;
  for(int i=0; i<sDegree+1; i++)
    delete [] c_x[i]; 
  for(int i=0; i<tDegree+1; i++)
    delete [] c_y[i];
  delete [] c_x;
  delete [] c_y;
  delete [] R;
  delete [] dR_dx;
  delete [] dR_dy;
  delete [] d2R_dx;
  delete [] d2R_dy;
  delete [] d2R_dxdy;
  std::cout<<"CLEAN: Element "<<en_index<<" is deleted"<<std::endl;
};

/* Element shape function routine */
void QuadElement::ElemShapeRoutine(double qua_x, double qua_y)
{
  int ii, jj;
  
  /* Initializing the Bernstein and B-spline data structure */
  
  double * Ns = new double [sDegree+1];
  double * Nt = new double [tDegree+1];
  double * dNs_ds = new double [sDegree+1];
  double * dNt_dt = new double [tDegree+1];
  double * d2Ns_ds = new double [sDegree+1];
  double * d2Nt_dt = new double [tDegree+1];

  for(ii=0; ii<sDegree+1; ii++)
  {
    Ns[ii]      = 0.0;
    dNs_ds[ii]  = 0.0;
    d2Ns_ds[ii] = 0.0;
  }

  for(ii=0;ii<tDegree+1;ii++)
  {
    Nt[ii]      = 0.0;
    dNt_dt[ii]  = 0.0;
    d2Nt_dt[ii] = 0.0;
  }

  double w           = 0.0;
  double dw_dxi      = 0.0;
  double dw_deta     = 0.0;
  double d2w_dxi    = 0.0;
  double d2w_deta   = 0.0;
  double d2w_dxideta = 0.0;

  /* Tensor product the Ns and Nt basis */
  double * N           = new double [enLocBas];
  double * dN_dxi      = new double [enLocBas];
  double * dN_deta     = new double [enLocBas];
  double * d2N_dxi     = new double [enLocBas];
  double * d2N_deta    = new double [enLocBas];
  double * d2N_dxideta = new double [enLocBas];

  /* NURBS basis */
  double * dR_dxi      = new double [enLocBas];
  double * dR_deta     = new double [enLocBas];
  double * d2R_dxi     = new double [enLocBas];
  double * d2R_deta    = new double [enLocBas];
  double * d2R_dxideta = new double [enLocBas];

  
  for( ii=0; ii<enLocBas; ii++ )
  {
    N[ii]           = 0.0;
    dN_dxi[ii]      = 0.0;
    dN_deta[ii]     = 0.0;
    d2N_dxi[ii]     = 0.0;
    d2N_deta[ii]    = 0.0;
    d2N_dxideta[ii] = 0.0;
    
    R[ii]           = 0.0;
    dR_dxi[ii]      = 0.0;
    dR_deta[ii]     = 0.0;
    d2R_dxi[ii]     = 0.0;
    d2R_deta[ii]    = 0.0;
    d2R_dxideta[ii] = 0.0;
    dR_dx[ii]       = 0.0;
    dR_dy[ii]       = 0.0;
    d2R_dx[ii]      = 0.0;
    d2R_dy[ii]      = 0.0;
    d2R_dxdy[ii]    = 0.0;
  }
  
  
  /* Geometry */
  double dx_dxi = 0.0;
  double dx_deta = 0.0;
  double dy_dxi = 0.0;
  double dy_deta = 0.0;

  double d2x_dxi = 0.0;
  double d2x_deta = 0.0;
  double d2x_dxideta = 0.0;
  double d2y_dxi = 0.0;
  double d2y_deta = 0.0;
  double d2y_dxideta = 0.0;
  
  Jacobi = 0.0;

  /* ==============================
   * 2. Compute the Bernstein polynomial and B-spline
   *    using extraction operator
   * ============================== */
  
  BernsteinBasis Bs(sDegree, qua_x);
  BernsteinBasis Bt(tDegree, qua_y);

  for( ii=0; ii<sDegree+1; ii++ )
  {
    for( jj=0; jj<sDegree+1; jj++ )
    {
      Ns[ii]      = Ns[ii]      + c_x[ii][jj] * Bs.der0[jj];
      dNs_ds[ii]  = dNs_ds[ii]  + c_x[ii][jj] * Bs.der1[jj]/en_hx;
      d2Ns_ds[ii] = d2Ns_ds[ii] + c_x[ii][jj] * Bs.der2[jj]/(en_hx*en_hx);
    }
  }

  for( ii=0; ii<tDegree+1; ii++ )
  {
    for( jj=0; jj<tDegree+1; jj++ )
    {
      Nt[ii]      = Nt[ii]      + c_y[ii][jj] * Bt.der0[jj];
      dNt_dt[ii]  = dNt_dt[ii]  + c_y[ii][jj] * Bt.der1[jj]/en_hy;
      d2Nt_dt[ii] = d2Nt_dt[ii] + c_y[ii][jj] * Bt.der2[jj]/(en_hy*en_hy);
    }
  }

  for( jj=0; jj<tDegree+1; jj++ )
  {
    for( ii=0; ii<sDegree+1; ii++)
    {
      int index = (sDegree+1) * jj + ii;

      N[index] = Ns[ii] * Nt[jj] * ctrl_w[index];
      w = w + N[index];
      
      dN_dxi[index] = dNs_ds[ii] * Nt[jj] * ctrl_w[index];
      dw_dxi = dw_dxi + dN_dxi[index];

      dN_deta[index] = Ns[ii] * dNt_dt[jj] * ctrl_w[index];
      dw_deta = dw_deta + dN_deta[index];

      d2N_dxi[index] = d2Ns_ds[ii] * Nt[jj] * ctrl_w[index];
      d2w_dxi = d2w_dxi + d2N_dxi[index];

      d2N_deta[index] = Ns[ii] * d2Nt_dt[jj] * ctrl_w[index];
      d2w_deta = d2w_deta + d2N_deta[index];

      d2N_dxideta[index] = dNs_ds[ii] * dNt_dt[jj] * ctrl_w[index];
      d2w_dxideta = d2w_dxideta + d2N_dxideta[index];
    }
  }

  /* Devide by denominators to complete the computation of the function */
  for( ii = 0; ii<enLocBas; ii++ )
  {
    R[ii] = N[ii] / w;
    
    dR_dxi[ii]  = ( dN_dxi[ii]  - R[ii] * dw_dxi ) / w;
    dR_deta[ii] = ( dN_deta[ii] - R[ii] * dw_deta ) / w;
    
    d2R_dxi[ii] = ( d2N_dxi[ii] - R[ii] * d2w_dxi - 
                       2*dR_dxi[ii] * dw_dxi ) / w;

    d2R_deta[ii] = ( d2N_deta[ii] - R[ii] * d2w_deta 
                        - 2*dR_deta[ii] * dw_deta ) / w;

    d2R_dxideta[ii] = ( d2N_dxideta[ii] - dR_deta[ii]*dw_dxi 
                          - R[ii] * d2w_dxideta - dR_dxi[ii]*dw_deta ) / w;
  }

  /* compute the mapping from the parameter space to physical space */
  for( ii = 0; ii<enLocBas; ii++)
  {
    dx_dxi  = dx_dxi  + ctrl_x[ii] * dR_dxi[ii];
    dx_deta = dx_deta + ctrl_x[ii] * dR_deta[ii];

    dy_dxi  = dy_dxi  + ctrl_y[ii] * dR_dxi[ii];
    dy_deta = dy_deta + ctrl_y[ii] * dR_deta[ii];

    d2x_dxi  = d2x_dxi  + ctrl_x[ii] * d2R_dxi[ii];
    d2x_deta = d2x_deta + ctrl_x[ii] * d2R_deta[ii];
    d2x_dxideta = d2x_dxideta + ctrl_x[ii] * d2R_dxideta[ii];

    d2y_dxi  = d2y_dxi  + ctrl_y[ii] * d2R_dxi[ii];
    d2y_deta = d2y_deta + ctrl_y[ii] * d2R_deta[ii];
    d2y_dxideta = d2y_dxideta + ctrl_y[ii] * d2R_dxideta[ii];
  }

  Jacobi = dx_dxi * dy_deta - dx_deta * dy_dxi;

  double dxi_dx = dy_deta / Jacobi;
  double dxi_dy = (-1.0) * dx_deta / Jacobi;
  double deta_dx = (-1.0) * dy_dxi / Jacobi;
  double deta_dy = dx_dxi / Jacobi;

  /* Compute the derivative w.r.t. the physical coordinate */
  for( ii = 0; ii<enLocBas; ii++ )
  {
    dR_dx[ii] = dR_dx[ii] + dR_dxi[ii] * dxi_dx + dR_deta[ii] * deta_dx;
    dR_dy[ii] = dR_dy[ii] + dR_dxi[ii] * dxi_dy + dR_deta[ii] * deta_dy;
  }

  Jacobi = fabs(Jacobi);
  Jacobi = Jacobi * en_hx * en_hy;

  /* To compute the second order derivatives, we set up a linear 
   * system of the form Ax = b
   * x is the vector of d2R/dx d2R/dxdy d2R/dy */

  double f1, f2, f3;
  double sol[3];

  for(int ii=0; ii<enLocBas; ii++)
  {
    // set up right hand side vector
    f1 = d2R_dxi[ii] - dR_dx[ii]*d2x_dxi - dR_dy[ii]*d2y_dxi;
    f2 = d2R_dxideta[ii] - dR_dx[ii]*d2x_dxideta - dR_dy[ii]*d2y_dxideta;
    f3 = d2R_deta[ii] - dR_dx[ii]*d2x_deta - dR_dy[ii]*d2y_deta;

    LinearSolver(dx_dxi, dy_dxi, dx_deta, dy_deta, f1, f2, f3, 
                d2R_dx[ii], d2R_dxdy[ii], d2R_dy[ii]);
  }
};

void QuadElement::LinearSolver( double a, double b, double c, double d, 
                                double f1, double f2, double f3, 
                                double &sol1, double &sol2, double &sol3 )
{
  /* This is the special routine for us to solve the equation associated
   * with the 2nd derivative w.r.t. physical coordinate
   * here a = dx_dxi, b = dy_dxi, c = dx_deta, d = dy_deta
   * f1 f2 f3 are the componenets of right hand vector:
   * f1 = d2R_dxi - dR_dx * d2x_dxi - dR_dy * d2y_dxi
   * f2 = d2R_dxideta - dR_dx * d2x_dxideta - dR_dy * d2y_dxideta
   * f3 = d2R_deta - dR_dx * d2x_deta - dR_dy * d2y_deta
   *
   * sol is the solution of A*sol = f
   *
   * A = [ a^2, 2*a*b, b^2;
   *       a*c, a*d+b*c, b*d;
   *       c^2, 2*c*d, d^2]                                        */

  sol1 = d*d*f1 - 2*b*d*f2 + b*b*f3;
  sol2 = -c*d*f1 + (a*d+b*c)*f2 - a*b*f3;
  sol3 = c*c*f1 - 2*a*c*f2 + a*a*f3;
  
  double bottom = (a*d-b*c)*(a*d-b*c);

  sol1 = sol1 / bottom;
  sol2 = sol2 / bottom;
  sol3 = sol3 / bottom;

};
